#!/usr/bin/env python
######################################################################################
#
#   This module is attempting to be a Python port of nw_rgb_make. It's still in 
#       development, but the basic use cases should work. The two types of FITS
#       files accepted for creation of an RGB image are 3 separate band pass images,
#       passed by using the -r, -g, and -b flags to control which color they get 
#       mapped to, or by the -f flag and passing a FITS file with NAXIS=3 in the 
#       primary HDU - where the 3rd dimension is color.
#
######################################################################################

__author__ = "Adrian Price-Whelan <adrn@nyu.edu>"

import sys, os
import Image
import numpy as np
import pyfits as pf
import APWRaw

def main():
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option("-r", "--r-band", dest="r", default="",
        help="A path to a FITS file containing image data to be mapped to 'R'")
    parser.add_option("-g", "--g-band", dest="g", default="",
        help="A path to a FITS file containing image data to be mapped to 'G'")
    parser.add_option("-b", "--b-band", dest="b", default="",
        help="A path to a FITS file containing image data to be mapped to 'B'")
    parser.add_option("-f", "--file", dest="file", default="",
        help="A path to a FITS file containing 3-dimensional image data, where the 3rd dimension is Color")
    parser.add_option("-o", "--out-file", dest="outFile", default="",
        help="If specified, it saves the output image into the specified filename. If not specified, it simply displays the image as a bitmap")
    parser.add_option("-B", "--Beta", dest="beta", default=3.0, type=float,
        help="The nonlinearity factor for the arcsinh stretch")
    parser.add_option("-i", "--invert", dest="invert", default=False, action="store_true",
        help="Invert the image")
        
    # This feature currently doesn't do anything, I hope to fix in future versions..
    parser.add_option("-s", "--sky-subtract", dest="skySub", default=False, action="store_true",
        help="")
    
    (options, args) = parser.parse_args()
    
    if options.file == "":
        if not os.path.exists(options.r): raise IOError("File: %s not found!" % options.r)
        if not os.path.exists(options.g): raise IOError("File: %s not found!" % options.g)
        if not os.path.exists(options.b): raise IOError("File: %s not found!" % options.b)
        
        r = pf.open(options.r)[0].data
        g = pf.open(options.g)[0].data
        b = pf.open(options.b)[0].data
        
    else:
        if not os.path.exists(options.file): raise IOError("File: %s not found!" % options.file)
        path,ext = os.path.splitext(options.file)
        ext = ext.lower()
        if ext == ".fits" or ext == ".fit":
            r,g,b = APWRaw.fitsToBands(options.file)
        else:
            print "Currently only supports FITS files!"
            sys.exit(0)
            """
            im = Image.open(options.file)
            data = np.asarray(im, dtype=float)
            maxrgb = np.max(data, axis=2)
            r = APWRaw.normalize(data[:,:,0])
            g = APWRaw.normalize(data[:,:,1])
            b = APWRaw.normalize(data[:,:,2])
            """
    
    if options.skySub:
        R,G,B = APWRaw.luptonScale(APWRaw.skySubtract(r), \
                               APWRaw.skySubtract(g), \
                               APWRaw.skySubtract(b), \
                               beta=options.beta)
    else:
        R,G,B = APWRaw.luptonScale(r, \
                                   g, \
                                   b, \
                                   beta=options.beta,\
                                   scales=[3.,3.,3.])
    
    imData = np.zeros((R.shape[0],R.shape[1],3), dtype=np.uint8)
    imData[:,:,0] = (R*255.0).clip(0,255).astype(np.uint8)
    imData[:,:,1] = (G*255.0).clip(0,255).astype(np.uint8)
    imData[:,:,2] = (B*255.0).clip(0,255).astype(np.uint8)
    
    if options.invert:
        imData = 255 - imData
    
    img = Image.fromarray(imData, mode="RGB")
    if options.outFile == "":
        img.show()
    else:
        img.save(options.outFile)

def plotScaling():
    import matplotlib.pyplot as plt
    
    x = np.arange(0.0, 1.0, 0.005)
    y1 = APWRaw.apw_arcsinh(x, Q=0.001, alpha=10.0)
    y2 = APWRaw.apw_arcsinh(x, Q=0.001, alpha=1.0)
    y3 = APWRaw.apw_arcsinh(x, Q=10.0, alpha=10.0)
    y4 = APWRaw.apw_arcsinh(x, Q=10.0, alpha=1.0)
    
    plt.plot(x, y1, 'r-')
    plt.plot(x, y2, 'g-')
    plt.plot(x, y3, 'b-')
    plt.plot(x, y4, 'c-')
    plt.show()
    

if __name__ == "__main__":
    #plotScaling()
    main()
    # Test cases:
    # python apw_rgb_make.py -r ~/Downloads/frame-u-000752-1-0331.fits -g ~/Downloads/frame-g-000752-1-0331.fits -b ~/Downloads/frame-r-000752-1-0331.fits -B 3.0 -o sdss_test_apw.png
    # python apw_rgb_make.py -f DSC_2431.png